"""
# multigrid.jl  
# Multigrid methods and preconditioning for Spectral Element Methods

This module implements geometric and algebraic multigrid methods optimized for high-order
spectral element discretizations, including p-multigrid (polynomial order coarsening),
h-multigrid (element coarsening), and hybrid hp-multigrid strategies.
"""

using LinearAlgebra
using SparseArrays
using Printf

"""
    MultigridLevel

Represents one level in the multigrid hierarchy.

# Fields
- `operators::Any`: SEM operators at this level (mass, stiffness, etc.)
- `grid::Any`: Grid information at this level
- `polynomial_order::Int`: Polynomial order at this level
- `n_elements::Int`: Number of elements at this level
- `restriction::AbstractMatrix`: Restriction operator to coarse level
- `prolongation::AbstractMatrix`: Prolongation operator from coarse level
- `system_matrix::AbstractMatrix`: System matrix at this level
- `smoother::Any`: Smoother configuration for this level
"""
struct MultigridLevel
    operators::Any
    grid::Any
    polynomial_order::Int
    n_elements::Int
    restriction::Union{AbstractMatrix, Nothing}
    prolongation::Union{AbstractMatrix, Nothing}  
    system_matrix::Union{AbstractMatrix, Nothing}
    smoother::Any
end

"""
    MultigridPreconditioner

Multigrid preconditioner for spectral element systems.

# Fields
- `levels::Vector{MultigridLevel}`: Hierarchy of multigrid levels
- `smoother::Symbol`: Smoothing method (`:jacobi`, `:gauss_seidel`, `:sor`, `:chebyshev`)
- `cycle_type::Symbol`: Multigrid cycle type (`:V`, `:W`, `:F`, `:FMG`)
- `pre_smooth::Int`: Pre-smoothing iterations
- `post_smooth::Int`: Post-smoothing iterations
- `coarse_solver::Symbol`: Coarse grid solver (`:direct`, `:iterative`)
- `omega::Float64`: Relaxation parameter for smoothers
- `chebyshev_degree::Int`: Degree for Chebyshev smoothers
- `max_coarse_size::Int`: Maximum size for coarse grid direct solve
- `verbose::Bool`: Enable verbose output
"""
Base.@kwdef struct MultigridPreconditioner
    levels::Vector{MultigridLevel}
    smoother::Symbol = :jacobi
    cycle_type::Symbol = :V
    pre_smooth::Int = 2
    post_smooth::Int = 2
    coarse_solver::Symbol = :direct
    omega::Float64 = 0.7
    chebyshev_degree::Int = 3
    max_coarse_size::Int = 1000
    verbose::Bool = false
end

"""
    create_multigrid_hierarchy(sem_operators, multidomain, max_levels::Int=4;
                              coarsening_strategy::Symbol=:p_multigrid,
                              min_polynomial_order::Int=2) -> Vector{MultigridLevel}

Create multigrid hierarchy for spectral element system.

# Arguments
- `sem_operators`: SEM operators at finest level
- `multidomain`: Multi-domain setup
- `max_levels::Int`: Maximum number of levels in hierarchy
- `coarsening_strategy::Symbol`: Strategy (`:p_multigrid`, `:h_multigrid`, `:hp_multigrid`)
- `min_polynomial_order::Int`: Minimum polynomial order for p-coarsening

# Returns
- `Vector{MultigridLevel}`: Complete multigrid hierarchy
"""
function create_multigrid_hierarchy(sem_operators, multidomain, max_levels::Int=4;
                                   coarsening_strategy::Symbol=:p_multigrid,
                                   min_polynomial_order::Int=2)
    
    println("🏗️  Creating multigrid hierarchy with $max_levels levels")
    println("   Strategy: $coarsening_strategy")
    
    levels = MultigridLevel[]
    
    # Create finest level
    finest_level = MultigridLevel(
        sem_operators,
        multidomain,
        multidomain.n,
        multidomain.index_sum,
        nothing,  # No restriction from finer level
        nothing,  # No prolongation to finer level
        nothing,  # System matrix created later
        create_smoother_config(:jacobi, 0.7)
    )
    push!(levels, finest_level)
    
    # Create coarser levels
    current_n = multidomain.n
    current_operators = sem_operators
    current_domain = multidomain
    
    for level in 2:max_levels
        if coarsening_strategy == :p_multigrid
            # p-multigrid: reduce polynomial order
            new_n = max(min_polynomial_order, current_n - 1)
            if new_n >= current_n
                println("   Stopping p-coarsening at level $level (minimum order reached)")
                break
            end
            
            # Create coarse SEM operators
            coarse_operators = create_coarse_sem_operators(current_operators, new_n)
            coarse_domain = create_coarse_domain(current_domain, new_n, :p_coarsen)
            
            # Create intergrid operators
            restriction_op = create_p_restriction_operator(current_n, new_n)
            prolongation_op = create_p_prolongation_operator(new_n, current_n)
            
        elseif coarsening_strategy == :h_multigrid
            # h-multigrid: reduce number of elements
            if level > 3
                println("   Stopping h-coarsening at level $level (implementation limit)")
                break
            end
            
            coarse_operators = create_h_coarse_operators(current_operators)
            coarse_domain = create_coarse_domain(current_domain, current_n, :h_coarsen)
            
            restriction_op = create_h_restriction_operator(current_domain, coarse_domain)
            prolongation_op = create_h_prolongation_operator(coarse_domain, current_domain)
            new_n = current_n  # Polynomial order unchanged
            
        else  # :hp_multigrid
            # Combined hp-coarsening strategy
            if current_n > min_polynomial_order
                new_n = max(min_polynomial_order, current_n - 1)
                coarse_operators = create_coarse_sem_operators(current_operators, new_n)
                coarse_domain = create_coarse_domain(current_domain, new_n, :p_coarsen)
                restriction_op = create_p_restriction_operator(current_n, new_n)
                prolongation_op = create_p_prolongation_operator(new_n, current_n)
            else
                println("   Stopping hp-coarsening at level $level (minimum order reached)")
                break
            end
        end
        
        # Create multigrid level
        coarse_level = MultigridLevel(
            coarse_operators,
            coarse_domain,
            new_n,
            coarse_domain.index_sum,
            restriction_op,
            prolongation_op,
            nothing,  # System matrix created later
            create_smoother_config(:jacobi, 0.7)
        )
        
        push!(levels, coarse_level)
        
        println("   Level $level: N=$(new_n), Elements=$(coarse_domain.index_sum)")
        
        # Update for next level
        current_n = new_n
        current_operators = coarse_operators
        current_domain = coarse_domain
    end
    
    println("   Created $(length(levels)) levels total")
    return levels
end

"""
    mg_cycle!(x::AbstractArray, b::AbstractArray, mg::MultigridPreconditioner, level::Int=1)

Perform one multigrid cycle (V, W, or F-cycle).
"""
function mg_cycle!(x::AbstractArray, b::AbstractArray, mg::MultigridPreconditioner, level::Int=1)
    
    if level == length(mg.levels)
        # Coarse grid solve
        solve_coarse_system!(x, b, mg.levels[level], mg)
        return
    end
    
    current_level = mg.levels[level]
    coarse_level = mg.levels[level + 1]
    
    # Pre-smoothing
    for i in 1:mg.pre_smooth
        smooth!(x, b, current_level, mg)
    end
    
    # Compute residual
    residual = b - current_level.system_matrix * x
    
    # Restrict residual to coarse grid
    coarse_residual = coarse_level.restriction * residual
    coarse_correction = zeros(length(coarse_residual))
    
    # Recursive coarse grid correction
    if mg.cycle_type == :V
        mg_cycle!(coarse_correction, coarse_residual, mg, level + 1)
    elseif mg.cycle_type == :W
        mg_cycle!(coarse_correction, coarse_residual, mg, level + 1)
        mg_cycle!(coarse_correction, coarse_residual, mg, level + 1)  # Second call
    elseif mg.cycle_type == :F
        # F-cycle implementation (nested V-cycles)
        for f_iter in 1:level
            mg_cycle!(coarse_correction, coarse_residual, mg, level + 1)
        end
    end
    
    # Prolongate correction and apply
    correction = coarse_level.prolongation * coarse_correction
    x .+= correction
    
    # Post-smoothing
    for i in 1:mg.post_smooth
        smooth!(x, b, current_level, mg)
    end
end

"""
    smooth!(x::AbstractArray, b::AbstractArray, level::MultigridLevel, mg::MultigridPreconditioner)

Apply smoother at given multigrid level.
"""
function smooth!(x::AbstractArray, b::AbstractArray, level::MultigridLevel, mg::MultigridPreconditioner)
    
    A = level.system_matrix
    n = size(A, 1)
    
    if mg.smoother == :jacobi
        # Weighted Jacobi smoother
        D_inv = Diagonal([A[i,i] != 0 ? 1/A[i,i] : 1.0 for i in 1:n])
        residual = b - A * x
        x .+= mg.omega * D_inv * residual
        
    elseif mg.smoother == :gauss_seidel
        # Gauss-Seidel smoother
        for i in 1:n
            if A[i,i] != 0
                sigma = dot(A[i,:], x) - A[i,i] * x[i]
                x[i] = (b[i] - sigma) / A[i,i]
            end
        end
        
    elseif mg.smoother == :sor
        # Successive over-relaxation
        for i in 1:n
            if A[i,i] != 0
                sigma = dot(A[i,1:i-1], x[1:i-1]) + dot(A[i,i+1:n], x[i+1:n])
                x[i] = (1 - mg.omega) * x[i] + mg.omega * (b[i] - sigma) / A[i,i]
            end
        end
        
    elseif mg.smoother == :chebyshev
        # Chebyshev polynomial smoother
        apply_chebyshev_smoother!(x, b, A, mg.chebyshev_degree)
        
    else
        @warn "Unknown smoother type: $(mg.smoother), using Jacobi"
        D_inv = Diagonal([A[i,i] != 0 ? 1/A[i,i] : 1.0 for i in 1:n])
        residual = b - A * x
        x .+= mg.omega * D_inv * residual
    end
end

"""
    solve_coarse_system!(x::AbstractArray, b::AbstractArray, coarse_level::MultigridLevel, mg::MultigridPreconditioner)

Solve coarse grid system.
"""
function solve_coarse_system!(x::AbstractArray, b::AbstractArray, coarse_level::MultigridLevel, mg::MultigridPreconditioner)
    
    A = coarse_level.system_matrix
    
    if mg.coarse_solver == :direct && size(A, 1) <= mg.max_coarse_size
        # Direct solve for small systems
        try
            x .= A \ b
        catch
            @warn "Direct coarse solve failed, using iterative"
            x, _ = gmres(A, b; abstol=1e-6, maxiter=min(100, size(A,1)))
        end
    else
        # Iterative solve
        x, info = gmres(A, b; abstol=1e-8, maxiter=min(200, size(A,1)))
        if !info.converged && mg.verbose
            @warn "Coarse grid solve did not converge (residual: $(@sprintf("%.2e", norm(A*x - b))))"
        end
    end
end

"""
    create_p_restriction_operator(fine_order::Int, coarse_order::Int) -> AbstractMatrix

Create restriction operator for p-multigrid (polynomial order coarsening).
"""
function create_p_restriction_operator(fine_order::Int, coarse_order::Int)
    
    if coarse_order >= fine_order
        return I  # Identity if no coarsening
    end
    
    # Create restriction based on Legendre polynomial projections
    # This is a simplified implementation - full version would use proper SEM basis functions
    
    fine_size = (fine_order + 1)^2  # 2D case
    coarse_size = (coarse_order + 1)^2
    
    # Simple truncation restriction (could be improved with L2 projection)
    restriction = zeros(coarse_size, fine_size)
    
    fine_1d = fine_order + 1
    coarse_1d = coarse_order + 1
    
    for j in 1:coarse_1d
        for i in 1:coarse_1d
            coarse_idx = (j-1) * coarse_1d + i
            fine_idx = (j-1) * fine_1d + i
            if fine_idx <= fine_size
                restriction[coarse_idx, fine_idx] = 1.0
            end
        end
    end
    
    return sparse(restriction)
end

"""
    create_p_prolongation_operator(coarse_order::Int, fine_order::Int) -> AbstractMatrix

Create prolongation operator for p-multigrid.
"""
function create_p_prolongation_operator(coarse_order::Int, fine_order::Int)
    
    if fine_order <= coarse_order
        return I  # Identity if no refinement
    end
    
    # Create prolongation (transpose of restriction with appropriate scaling)
    restriction = create_p_restriction_operator(fine_order, coarse_order)
    
    # Simple prolongation (could be improved with interpolation)
    return restriction'
end

"""
    create_h_restriction_operator(fine_domain, coarse_domain) -> AbstractMatrix

Create restriction operator for h-multigrid (element coarsening).
"""
function create_h_restriction_operator(fine_domain, coarse_domain)
    # Placeholder implementation for h-multigrid
    # Would require careful handling of element-to-element mappings
    
    fine_size = fine_domain.index_sum * (fine_domain.n + 1)^fine_domain.dim
    coarse_size = coarse_domain.index_sum * (coarse_domain.n + 1)^coarse_domain.dim
    
    # Simple restriction (would be more sophisticated in practice)
    return sparse(I, min(coarse_size, fine_size), fine_size)
end

"""
    create_h_prolongation_operator(coarse_domain, fine_domain) -> AbstractMatrix

Create prolongation operator for h-multigrid.
"""
function create_h_prolongation_operator(coarse_domain, fine_domain)
    restriction = create_h_restriction_operator(fine_domain, coarse_domain)
    return restriction'
end

"""
    create_coarse_sem_operators(fine_operators, coarse_order::Int)

Create SEM operators at coarse polynomial order.
"""
function create_coarse_sem_operators(fine_operators, coarse_order::Int)
    # Placeholder - would create actual coarse SEM operators
    # This would involve re-computing mass matrices, stiffness matrices, etc.
    # at the coarse polynomial order
    
    return fine_operators  # Placeholder
end

"""
    create_coarse_domain(fine_domain, coarse_order::Int, strategy::Symbol)

Create coarse domain for multigrid level.
"""
function create_coarse_domain(fine_domain, coarse_order::Int, strategy::Symbol)
    if strategy == :p_coarsen
        # Create new domain with reduced polynomial order
        return typeof(fine_domain)(
            coarse_order,
            fine_domain.n_block,
            fine_domain.dim,
            coarse_order + 1,
            fine_domain.length,
            fine_domain.time_param,
            fine_domain.matrix_dist,
            fine_domain.g,
            fine_domain.gi,
            fine_domain.index_sum,
            fine_domain.adaptive_refinement,
            fine_domain.refinement_map,
            fine_domain.corner_points
        )
    else
        # For h-coarsening, would modify element structure
        return fine_domain  # Placeholder
    end
end

"""
    create_smoother_config(smoother_type::Symbol, omega::Float64)

Create smoother configuration object.
"""
function create_smoother_config(smoother_type::Symbol, omega::Float64)
    return (type=smoother_type, omega=omega)
end

"""
    apply_chebyshev_smoother!(x::AbstractArray, b::AbstractArray, A::AbstractMatrix, degree::Int)

Apply Chebyshev polynomial smoother.
"""
function apply_chebyshev_smoother!(x::AbstractArray, b::AbstractArray, A::AbstractMatrix, degree::Int)
    
    # Estimate spectral radius
    lambda_max = estimate_spectral_radius(A)
    lambda_min = lambda_max / 30.0  # Typical ratio for SEM systems
    
    # Chebyshev parameters
    center = (lambda_max + lambda_min) / 2
    radius = (lambda_max - lambda_min) / 2
    
    # Chebyshev iteration
    residual = b - A * x
    p_old = zeros(size(x))
    p = residual / center
    x .+= p
    
    for k in 2:degree
        beta = (radius * 0.5)^2 / center
        if k == 2
            beta = beta / 2
        end
        
        residual = b - A * x
        p_new = residual / center + beta * p - beta * p_old
        x .+= p_new
        
        p_old = p
        p = p_new
    end
end

"""
    estimate_spectral_radius(A::AbstractMatrix) -> Float64

Estimate spectral radius of matrix using power iteration.
"""
function estimate_spectral_radius(A::AbstractMatrix)
    n = size(A, 1)
    x = randn(n)
    x /= norm(x)
    
    lambda_old = 0.0
    for i in 1:10  # Few iterations for estimate
        y = A * x
        lambda_new = dot(x, y)
        x = y / norm(y)
        
        if abs(lambda_new - lambda_old) < 1e-3 * abs(lambda_new)
            break
        end
        lambda_old = lambda_new
    end
    
    return abs(lambda_old)
end

"""
    setup_multigrid_preconditioner!(mg::MultigridPreconditioner, system_matrices::Vector{AbstractMatrix})

Set up system matrices for all multigrid levels.
"""
function setup_multigrid_preconditioner!(mg::MultigridPreconditioner, system_matrices::Vector{AbstractMatrix})
    
    if length(system_matrices) != length(mg.levels)
        error("Number of system matrices must match number of multigrid levels")
    end
    
    for (i, matrix) in enumerate(system_matrices)
        # Create new level with system matrix
        level = mg.levels[i]
        mg.levels[i] = MultigridLevel(
            level.operators,
            level.grid,
            level.polynomial_order,
            level.n_elements,
            level.restriction,
            level.prolongation,
            matrix,
            level.smoother
        )
    end
    
    if mg.verbose
        println("📋 Multigrid preconditioner setup complete")
        for (i, level) in enumerate(mg.levels)
            matrix_size = size(level.system_matrix, 1)
            println("   Level $i: $(matrix_size)×$(matrix_size) matrix, order=$(level.polynomial_order)")
        end
    end
end

"""
    apply_multigrid_preconditioner!(y::AbstractArray, x::AbstractArray, mg::MultigridPreconditioner)

Apply multigrid preconditioner: y = M^{-1} x
"""
function apply_multigrid_preconditioner!(y::AbstractArray, x::AbstractArray, mg::MultigridPreconditioner)
    # Initialize solution
    y .= 0
    
    # Apply multigrid cycle as preconditioner
    mg_cycle!(y, x, mg, 1)
end

"""
    create_sem_multigrid_preconditioner(sem_operators, multidomain, system_matrix::AbstractMatrix;
                                      max_levels::Int=4, cycle_type::Symbol=:V,
                                      smoother::Symbol=:jacobi, verbose::Bool=false) -> MultigridPreconditioner

Create complete multigrid preconditioner for SEM system.
"""
function create_sem_multigrid_preconditioner(sem_operators, multidomain, system_matrix::AbstractMatrix;
                                            max_levels::Int=4, cycle_type::Symbol=:V,
                                            smoother::Symbol=:jacobi, verbose::Bool=false)
    
    if verbose
        println("🔧 Creating SEM multigrid preconditioner")
        println("   System size: $(size(system_matrix))")
        println("   Cycle type: $cycle_type, Smoother: $smoother")
    end
    
    # Create multigrid hierarchy
    levels = create_multigrid_hierarchy(sem_operators, multidomain, max_levels)
    
    # Create preconditioner
    mg_preconditioner = MultigridPreconditioner(
        levels=levels,
        smoother=smoother,
        cycle_type=cycle_type,
        pre_smooth=2,
        post_smooth=2,
        coarse_solver=:direct,
        verbose=verbose
    )
    
    # Set up system matrices (for now, just finest level)
    # In practice, would compute coarse matrices via Galerkin projection: A_coarse = R * A_fine * P
    system_matrices = [system_matrix]
    
    for level in 2:length(levels)
        # Galerkin coarse grid operator: A_coarse = R * A_fine * P
        R = levels[level].restriction
        P = levels[level].prolongation
        A_fine = system_matrices[level-1]
        
        if R !== nothing && P !== nothing
            A_coarse = R * A_fine * P
            push!(system_matrices, A_coarse)
        else
            # Fallback: use same matrix (not recommended for production)
            push!(system_matrices, A_fine)
        end
    end
    
    setup_multigrid_preconditioner!(mg_preconditioner, system_matrices)
    
    return mg_preconditioner
end